year_from <- min(sweden_deaths_monthly$Year)

year_reg <- year_from + 10

year_max <- 2020

flu_years_hc <- c(1890, 1918, 2020)

flu_years_hc_affected <- c(seq(1890 + 1, 1890 + 10),
                           seq(1918 + 1, 1918 + 10))

flu_years <- c(1890, 1918, 1957, 1968, 1977, 2009, 2020)

flu_years_affected <- c(flu_years_hc_affected,
                        seq(1957 + 1, 1957 + 10),
                        seq(1968 + 1, 1968 + 10),
                        seq(1977 + 1, 1977 + 10),
                        seq(2009 + 1, 2009 + 10))

sweden_deaths_monthly %<>% 
  mutate(flu_year_hc = ifelse(Year %in% flu_years_hc, 1, 0),
         flu_year_hc_affected = ifelse(Year %in% flu_years_hc_affected, 1, 0),
         deaths_flu_year_hc = ifelse(Year %in% flu_years_hc, NA, deaths),
         flu_year = ifelse(Year %in% flu_years, 1, 0),
         flu_year_affected = ifelse(Year %in% flu_years_affected, 1, 0),
         deaths_flu_year = ifelse(Year %in% flu_years, NA, deaths),
         si_one = sin(2*pi*Month/12),
         si_two = sin(4*pi*Month/12),
         co_one = cos(2*pi*Month/12),
         co_two = cos(4*pi*Month/12)) 

sweden_deaths_monthly %<>% 
  arrange(Month, Year) %>% 
  mutate(median =             slide_dbl(deaths, 
                                        stats::median, na.rm = TRUE, 
                                        .before = 10, .after = -1, .complete = TRUE),
         median_flu_year_hc = slide_dbl(deaths_flu_year_hc, 
                                        stats::median, na.rm = TRUE,
                                        .before = 10, .after = -1, .complete = TRUE),
         median_flu_year =    slide_dbl(deaths_flu_year, 
                                        stats::median, na.rm = TRUE, 
                                        .before = 10, .after = -1, .complete = TRUE),
  ) %>% 
  arrange(Year, Month)

dcmp <- sweden_deaths_monthly %>%
  mutate(Time = row_number()) %>% 
  tsibble::as_tsibble(index = Time) %>% 
  fabletools::model(feasts::STL((deaths/pop)*100000 ~ season(window = Inf)))

# fabletools::components(dcmp)

fabletools::components(dcmp) %>% 
  autoplot()

p_load(doParallel, foreach)

registerDoParallel(cores = 4)

boot_pi <- function(model, pdata, n, p) {
  odata <- model$data
  lp <- (1 - p) / 2
  up <- 1 - lp
  seeds <- round(runif(n, 1, 1000), 0)
  boot_y <- foreach(i = 1:n, .combine = rbind) %dopar% {
    set.seed(seeds[i])
    bdata <- odata[sample(seq(nrow(odata)), size = nrow(odata), replace = TRUE), ]
    bpred <- predict(update(model, data = bdata), type = "response", newdata = pdata)
    rpois(length(bpred), lambda = bpred)
  }
  boot_ci <- t(apply(boot_y, 2, quantile, c(lp, up)))
  return(data.frame(pred = predict(model, newdata = pdata, type = "response"), 
                    lower = boot_ci[, 1], upper = boot_ci[, 2]))
}

for (YEAR in year_reg:year_max){
  
  print(paste("Analysing year", YEAR))
  
  reg_data <- sweden_deaths_monthly %>% 
    filter(Year >= YEAR - 10 & Year < YEAR) 
  
  pred_data <- sweden_deaths_monthly %>% 
    filter(Year == YEAR) %>%
    mutate(
      fit = NA_real_,
      lpi = NA_real_,
      upi = NA_real_,
      fit_flu_hc = NA_real_,
      lpi_flu_hc = NA_real_,
      upi_flu_hc = NA_real_,
      fit_flu = NA_real_,
      lpi_flu = NA_real_,
      upi_flu = NA_real_)
  
  # summary(reg_data$Year)
  # summary(pred_data$Year)
  
  # all data
  summary(monthly <- glm(deaths ~ Year + 
                           si_one + si_two + co_one + co_two, 
                         offset = log(pop),
                         data = reg_data, 
                         family = "poisson"))
  
  predict <- boot_pi(monthly, pred_data, 1000, 0.99)
  
  pred_data %<>%
    mutate(
      fit = predict$pred,
      lpi = predict$lower,
      upi = predict$upper,
      fit_flu_hc = predict$pred,
      lpi_flu_hc = predict$lower,
      upi_flu_hc = predict$upper,
      fit_flu = predict$pred,
      lpi_flu = predict$lower,
      upi_flu = predict$upper
    )
  
  # flu hc excluded
  if (YEAR %in% flu_years_hc_affected) {
    
    print(paste("    >> Extra analyses for hc flu year exclusions"))
    
    summary(monthly_flu_hc <- glm(deaths_flu_year_hc ~ Year + 
                                    si_one + si_two + co_one + co_two, 
                                  offset = log(pop),
                                  data = reg_data, 
                                  family = "poisson"))
    
    predict_flu_hc <- boot_pi(monthly_flu_hc, pred_data, 1000, 0.99)
    
    pred_data %<>%
      mutate(
        fit_flu_hc = predict_flu_hc$pred,
        lpi_flu_hc = predict_flu_hc$lower,
        upi_flu_hc = predict_flu_hc$upper
      )
  }
  
  # all flu excluded
  if (YEAR %in% flu_years_affected) {
    
    print(paste("    >> Extra analyses for flu year exclusions"))
    
    summary(monthly_flu <- glm(deaths_flu_year ~ Year + 
                                 si_one + si_two + co_one + co_two, 
                               offset = log(pop),
                               data = reg_data, 
                               family = "poisson"))
    
    predict_flu <- boot_pi(monthly_flu, pred_data, 1000, 0.99)
    
    pred_data %<>%
      mutate(
        fit_flu = predict_flu$pred,
        lpi_flu = predict_flu$lower,
        upi_flu = predict_flu$upper
      )
  } 
  
  if (YEAR == year_reg) {
    se_expected_deaths_monthly <- pred_data
  } else {
    se_expected_deaths_monthly <- bind_rows(se_expected_deaths_monthly, pred_data)
  }
}

rm(predict, predict_flu_hc, predict_flu, pred_data, reg_data, YEAR)
write_rds(se_expected_deaths_monthly, "data/se_expected_deaths_monthly.Rds")
haven::write_dta(se_expected_deaths_monthly, "data/se_expected_deaths_monthly.dta")